{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# INT8 Quantization with Post-training Optimization Tool (POT) in Simplified Mode tutorial\n",
    "This tutorial was adapted from https://github.com/openvinotoolkit/openvino_notebooks/tree/main/notebooks/114-quantization-simplified-mode. Here, we use OpenVINO APIs provided by BigDL Nano instead to simplify the original tutorial.\n",
    "\n",
    "This tutorial shows how to quantize a [ResNet20](https://github.com/chenyaofo/pytorch-cifar-models) image classification model, trained on [CIFAR10 ](http://pytorch.org/vision/main/generated/torchvision.datasets.CIFAR10.html) dataset, using the Simplified Mode of OpenVINO Post-Training Optimization Tool (POT).\n",
    "\n",
    "Simplified Mode is designed to make the data preparation step easier, before model optimization. The mode is represented by an implementation of the engine interface in the POT API. It enables reading data from an arbitrary folder specified by the user. Currently, Simplified Mode is available only for image data in PNG or JPEG formats, stored in a single folder.\n",
    "\n",
    "**Note:** This mode cannot be used with the accuracy-aware method. It is not possible to control accuracy after optimization using this mode. However, Simplified Mode can be useful for estimating performance improvements when optimizing models.\n",
    "\n",
    "This tutorial includes the following steps:\n",
    "\n",
    "- Downloading and saving the CIFAR10 dataset\n",
    "- Preparing the model for quantization\n",
    "- Compressing the prepared model\n",
    "- Measuring and comparing the performance of the original and quantized models\n",
    "- Demonstrating the use of the quantized model for image classification\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import warnings\n",
    "\n",
    "import torch\n",
    "from torchvision import transforms as T\n",
    "from torchvision.datasets import CIFAR10\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from bigdl.nano.pytorch import InferenceOptimizer\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "# Set the data and model directories\n",
    "MODEL_DIR = 'model'\n",
    "CALIB_DIR = 'calib'\n",
    "CIFAR_DIR = 'cifar'\n",
    "CALIB_SET_SIZE = 300\n",
    "MODEL_NAME = 'resnet20'\n",
    "\n",
    "os.makedirs(MODEL_DIR, exist_ok=True)\n",
    "os.makedirs(CALIB_DIR, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare the calibration dataset\n",
    "The following steps are required to prepare the calibration dataset:\n",
    "- Download CIFAR10 dataset from Torchvision.datasets repository\n",
    "- Save the selected number of elements from this dataset as .png images in a separate folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = T.Compose([T.ToTensor()])\n",
    "dataset = CIFAR10(root=CIFAR_DIR, train=False, transform=transform, download=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare the Model\n",
    "Model preparation includes the following steps:,\n",
    "- Download PyTorch model from Torchvision repository,\n",
    "- Use `bigdl.nano.pytorch.InferenceOptimizer.trace(accelerator='openvino')` to convert Pytorch model to OpenVINO Intermediate Representation (IR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = torch.hub.load(\"chenyaofo/pytorch-cifar-models\", \"cifar10_resnet20\", pretrained=True)\n",
    "dummy_input = torch.randn(1, 3, 32, 32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we convert this model into the OpenVINO IR using `InferenceOptimizer.trace`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ov_model = InferenceOptimizer.trace(model,\n",
    "                         accelerator='openvino',\n",
    "                         input_sample=dummy_input)\n",
    "InferenceOptimizer.save(ov_model, MODEL_DIR)\n",
    "\n",
    "ir_model_xml = Path(MODEL_DIR) / ov_model.status['xml_path']\n",
    "ir_model_bin = ir_model_xml.with_suffix('.bin')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compression stage\n",
    "Compress the model with `InferenceOptimizer.quantize` from either traced OpenVINO model(Option 1) or Pytorch model(Option 2)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader = DataLoader(dataset, batch_size=1)\n",
    "# Option1: compress the model from traced OpenVINO model\n",
    "optimized_model = InferenceOptimizer.quantize(ov_model,\n",
    "                                   precision='int8',\n",
    "                                   accelerator='openvino',\n",
    "                                   calib_data=dataloader)\n",
    "\n",
    "# Option2: compress the model from Pytorch model\n",
    "# optimized_model = InferenceOptimizer.quantize(model,\n",
    "#                                    precision='int8',\n",
    "#                                    accelerator='openvino',\n",
    "#                                    calib_data=dataloader)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare Performance of the Original and Quantized Models\n",
    "\n",
    "Finally, we will measure the inference performance of the FP32 and INT8 models. To do this, we use [Benchmark Tool](https://docs.openvino.ai/latest/openvino_inference_engine_tools_benchmark_tool_README.html) - an inference performance measurement tool for OpenVINO.\n",
    "\n",
    "**NOTE:** For more accurate performance, we recommended running benchmark_app in a terminal/command prompt after closing other applications. Run benchmark_app -m model.xml -d CPU to benchmark async inference on CPU for one minute. Change CPU to GPU to benchmark on GPU. Run benchmark_app --help to see an overview of all command line options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_model_path = Path('compressed/optimized')\n",
    "\n",
    "InferenceOptimizer.save(optimized_model, optimized_model_path)\n",
    "\n",
    "optimized_model_xml = optimized_model_path / optimized_model.status['xml_path']\n",
    "optimized_model_bin = optimized_model_path / optimized_model.status['weight_path']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference FP32 model (IR)\n",
    "!benchmark_app -m $ir_model_xml -d CPU -api async -b 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference INT8 model (IR)\n",
    "!benchmark_app -m $optimized_model_xml -d CPU -api async -b 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Demonstration of the results\n",
    "\n",
    "This section demonstrates how to use the compressed model by running the optimized model on a subset of images from the CIFAR10 dataset and shows predictions using the model.\n",
    "\n",
    "The first step is to load the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compiled_model = InferenceOptimizer.load(optimized_model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define all possible labels from CIFAR10\n",
    "labels_names = [\"airplane\", \"automobile\", \"bird\", \"cat\", \"deer\", \"dog\", \"frog\", \"horse\", \"ship\", \"truck\"]\n",
    "all_images = []\n",
    "all_labels = []\n",
    "\n",
    "# get all images and their labels \n",
    "for batch in dataset:\n",
    "    all_images.append(torch.unsqueeze(batch[0], 0))\n",
    "    all_labels.append(batch[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This section defines the function that shows the images and their labels using the indexes and two lists created in the previous step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_pictures(indexes: list, images=all_images, labels=all_labels):\n",
    "    \"\"\"Plot images with the specified indexes.\n",
    "    :param indexes: a list of indexes of images to be displayed.\n",
    "    :param images: a list of images from the dataset.\n",
    "    :param labels: a list of labels for each image.\n",
    "    \"\"\"\n",
    "    num_pics = len(indexes)\n",
    "    _, axarr = plt.subplots(1, num_pics)\n",
    "    for idx, im_idx in enumerate(indexes):\n",
    "        assert idx < 10000, 'Cannot get such index, there are only 10000'\n",
    "        pic = np.rollaxis(images[im_idx].squeeze().numpy(), 0, 3)\n",
    "        axarr[idx].imshow(pic)\n",
    "        axarr[idx].set_title(labels_names[labels[im_idx]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section we define a function that uses the optimized model to obtain predictions for the selected images:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def infer_on_images(net, indexes: list, images=all_images):\n",
    "    \"\"\" Inference model on a set of images.\n",
    "    :param net: model on which do inference\n",
    "    :param indexes: a list of indexes of images to infer on.\n",
    "    :param images: a list of images from the dataset.\n",
    "    \"\"\"\n",
    "    predicted_labels = []\n",
    "    for idx in indexes:\n",
    "        assert idx < 10000, 'Cannot get such index, there are only 10000'\n",
    "        result = net(images[idx].detach())\n",
    "        predicted_labels.append(result)\n",
    "    return predicted_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "indexes_to_infer = [0, 1, 2]  # to plot specify indexes\n",
    "\n",
    "plot_pictures(indexes_to_infer)\n",
    "\n",
    "results_quanized = infer_on_images(compiled_model, indexes_to_infer)\n",
    "\n",
    "print(f\"Image labels using the quantized model : {results_quanized}.\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "K5HPrY_d-7cV",
    "E01dMaR2_AFL",
    "qMnYsGo9_MA8",
    "L0tH9KdwtHhV"
   ],
   "name": "INT8 Quantization with POT in Simplified Mode tutorial",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.7.13 ('BigDL')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "31767e73ccf5faf58ae8bd08f02e245390c91328ab147d812430b8cd7a11fd87"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
